package first.group.ssh.packet;

import first.group.ssh.SSHConsts;
import first.group.ssh.encrypt.Encryption;
import first.group.ssh.hostkey.HostKeyAlgorithm;
import first.group.ssh.keyexchange.KeyExchangeMethod;
import first.group.ssh.mac.MacAlgorithm;
import first.group.ssh.packet.reader.BooleanReader;
import first.group.ssh.packet.reader.ByteArrayReader;
import first.group.ssh.packet.reader.NameListReader;
import first.group.ssh.packet.writer.BooleanWriter;
import first.group.ssh.packet.writer.NameListWriter;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.util.Arrays;
import java.util.List;
import java.util.stream.Collectors;

public final class KeyExchangeInitPacket extends RawPacket {

    private final byte[] payload;
    public final byte[] cookie;
    public final List<String> keyExchangeMethods;
    public final List<String> hostKeyAlgorithms;
    public final List<String> clientToServerEncryptions;
    public final List<String> serverToClientEncryptions;
    public final List<String> clientToServerMacs;
    public final List<String> serverToClientMacs;
    public final List<String> clientToServerCompressions;
    public final List<String> serverToClientCompressions;
    public final List<String> clientToServerLanguages;
    public final List<String> serverToClientLanguages;
    public final boolean firstKexPacketFollows;

    public KeyExchangeInitPacket(byte[] cookie) {
        this(
                cookie,
                KeyExchangeMethod.methods.stream().map(o -> o.getName()).collect(Collectors.toUnmodifiableList()),
                HostKeyAlgorithm.algos.stream().map(o -> o.getName()).collect(Collectors.toUnmodifiableList()),
                Encryption.encryptions.stream().map(o -> o.getName()).collect(Collectors.toUnmodifiableList()),
                Encryption.encryptions.stream().map(o -> o.getName()).collect(Collectors.toUnmodifiableList()),
                MacAlgorithm.algos.stream().map(o -> o.getName()).collect(Collectors.toUnmodifiableList()),
                MacAlgorithm.algos.stream().map(o -> o.getName()).collect(Collectors.toUnmodifiableList()),
                List.of("none"),
                List.of("none"),
                List.of(),
                List.of(),
                false
        );
    }
    
    public KeyExchangeInitPacket(
            byte[] cookie,
            List<String> keyExchangeMethods,
            List<String> hostKeyAlgorithms,
            List<String> clientToServerEncryptions,
            List<String> serverToClientEncryptions,
            List<String> clientToServerMacs,
            List<String> serverToClientMacs,
            List<String> clientToServerCompressions,
            List<String> serverToClientCompressions,
            List<String> clientToServerLanguages,
            List<String> serverToClientLanguages,
            boolean firstKexPacketFollows,
            byte[] payload
    ) {
        super(SSHConsts.SSH_MSG_KEXINIT);
        this.cookie = cookie;
        this.keyExchangeMethods = keyExchangeMethods;
        this.hostKeyAlgorithms = hostKeyAlgorithms;
        this.clientToServerEncryptions = clientToServerEncryptions;
        this.serverToClientEncryptions = serverToClientEncryptions;
        this.clientToServerMacs = clientToServerMacs;
        this.serverToClientMacs = serverToClientMacs;
        this.clientToServerCompressions = clientToServerCompressions;
        this.serverToClientCompressions = serverToClientCompressions;
        this.clientToServerLanguages = clientToServerLanguages;
        this.serverToClientLanguages = serverToClientLanguages;
        this.firstKexPacketFollows = firstKexPacketFollows;
        this.payload = payload;
    }

    private static byte[] writePayload(
            byte[] cookie,
            List<String> keyExchangeMethods,
            List<String> hostKeyAlgorithms,
            List<String> clientToServerEncryptions,
            List<String> serverToClientEncryptions,
            List<String> clientToServerMacs,
            List<String> serverToClientMacs,
            List<String> clientToServerCompressions,
            List<String> serverToClientCompressions,
            List<String> clientToServerLanguages,
            List<String> serverToClientLanguages,
            boolean firstKexPacketFollows
    ) {
        var buffer = ByteBuffer.allocate(4096);
        buffer.order(ByteOrder.BIG_ENDIAN);

        buffer.put(cookie);
        keyExchangeMethods.stream().collect(NameListWriter.of(buffer));
        hostKeyAlgorithms.stream().collect(NameListWriter.of(buffer));
        clientToServerEncryptions.stream().collect(NameListWriter.of(buffer));
        serverToClientEncryptions.stream().collect(NameListWriter.of(buffer));
        clientToServerMacs.stream().collect(NameListWriter.of(buffer));
        serverToClientMacs.stream().collect(NameListWriter.of(buffer));
        clientToServerCompressions.stream().collect(NameListWriter.of(buffer));
        serverToClientCompressions.stream().collect(NameListWriter.of(buffer));
        clientToServerLanguages.stream().collect(NameListWriter.of(buffer));
        serverToClientLanguages.stream().collect(NameListWriter.of(buffer));
        BooleanWriter.writeBoolean(buffer, firstKexPacketFollows);
        buffer.putInt(0);

        buffer.flip();

        var data = new byte[buffer.limit()];
        buffer.get(data);
        return data;
    }

    public KeyExchangeInitPacket(
            byte[] cookie,
            List<String> keyExchangeMethods,
            List<String> hostKeyAlgorithms,
            List<String> clientToServerEncryptions,
            List<String> serverToClientEncryptions,
            List<String> clientToServerMacs,
            List<String> serverToClientMacs,
            List<String> clientToServerCompressions,
            List<String> serverToClientCompressions,
            List<String> clientToServerLanguages,
            List<String> serverToClientLanguages,
            boolean firstKexPacketFollows
    ) {
        this(
                cookie,
                keyExchangeMethods,
                hostKeyAlgorithms,
                clientToServerEncryptions,
                serverToClientEncryptions,
                clientToServerMacs,
                serverToClientMacs,
                clientToServerCompressions,
                serverToClientCompressions,
                clientToServerLanguages,
                serverToClientLanguages,
                firstKexPacketFollows,
                writePayload(
                        cookie,
                        keyExchangeMethods,
                        hostKeyAlgorithms,
                        clientToServerEncryptions,
                        serverToClientEncryptions,
                        clientToServerMacs,
                        serverToClientMacs,
                        clientToServerCompressions,
                        serverToClientCompressions,
                        clientToServerLanguages,
                        serverToClientLanguages,
                        firstKexPacketFollows
                )
        );
    }

    @Override
    public byte[] getPayload() {
        return payload;
    }

    public static PacketParser<KeyExchangeInitPacket> PARSER = new Parser();

    private static class Parser extends PacketParser<KeyExchangeInitPacket> {

        @Override
        public KeyExchangeInitPacket parse(ByteBuffer buffer) throws IOException {
            var cookie = ByteArrayReader.readNameList(buffer, 16);
            if (cookie == null) {
                return null;
            }
            var keyExchangeMethods = NameListReader.readNameList(buffer);
            if (keyExchangeMethods == null) {
                return null;
            }
            var hostKeyAlgorithms = NameListReader.readNameList(buffer);
            if (hostKeyAlgorithms == null) {
                return null;
            }
            var clientToServerEncryptions = NameListReader.readNameList(buffer);
            if (clientToServerEncryptions == null) {
                return null;
            }
            var serverToClientEncryptions = NameListReader.readNameList(buffer);
            if (serverToClientEncryptions == null) {
                return null;
            }
            var clientToServerMacs = NameListReader.readNameList(buffer);
            if (clientToServerMacs == null) {
                return null;
            }
            var serverToClientMacs = NameListReader.readNameList(buffer);
            if (serverToClientMacs == null) {
                return null;
            }
            var clientToServerCompressions = NameListReader.readNameList(buffer);
            if (clientToServerCompressions == null) {
                return null;
            }
            var serverToClientCompressions = NameListReader.readNameList(buffer);
            if (serverToClientCompressions == null) {
                return null;
            }
            var clientToServerLanguages = NameListReader.readNameList(buffer);
            if (clientToServerLanguages == null) {
                return null;
            }
            var serverToClientLanguages = NameListReader.readNameList(buffer);
            if (serverToClientLanguages == null) {
                return null;
            }
            var firstKexPacketFollows = BooleanReader.readBoolean(buffer);
            if (firstKexPacketFollows == null) {
                return null;
            }
            if (!bufferRemains(buffer, 4)) {
                return null;
            }
            var reserved = buffer.getInt();
            if (reserved != 0) {
                throw new IOException("The reserved value should be 0, but get" + reserved);
            }
            return new KeyExchangeInitPacket(
                    cookie,
                    keyExchangeMethods,
                    hostKeyAlgorithms,
                    clientToServerEncryptions,
                    serverToClientEncryptions,
                    clientToServerMacs,
                    serverToClientMacs,
                    clientToServerCompressions,
                    serverToClientCompressions,
                    clientToServerLanguages,
                    serverToClientLanguages,
                    firstKexPacketFollows
            );
        }
    }

    @Override
    public void write(ByteBuffer buffer) {
        buffer.put(payload);
    }

    @Override
    public String toString() {
        return "KeyExchangeInitPacket{"
                + "cookie=" + Arrays.toString(cookie)
                + ", keyExchangeMethods=" + keyExchangeMethods
                + ", hostKeyAlgorithms=" + hostKeyAlgorithms
                + ", clientToServerEncryptions=" + clientToServerEncryptions
                + ", serverToClientEncryptions=" + serverToClientEncryptions
                + ", clientToServerMacs=" + clientToServerMacs
                + ", serverToClientMacs=" + serverToClientMacs
                + ", clientToServerCompressions=" + clientToServerCompressions
                + ", serverToClientCompressions=" + serverToClientCompressions
                + ", clientToServerLanguages=" + clientToServerLanguages
                + ", serverToClientLanguages=" + serverToClientLanguages
                + ", firstKexPacketFollows=" + firstKexPacketFollows + '}';
    }
}
